Revenge of "The Salary of AtCoder Inc."
/icons/hr.icon
問題
/icons/hr.icon
出典
ABC326 - E 1363 diff
/icons/hr.icon
問題文
AtCoder 社の社員である青木さんの今月の給料は、整数 $ N と長さ$ N の数列 $ A を用いて以下のように決められます。
まず、青木さんに $ 1 から $ N までの整数が等確率で出る $ N 面ダイスと変数 $ x=0 を渡します。
その後、以下の手順を終了まで繰り返します。
ダイスを $ 1 度振り、出た目を $ y とする。
もし $ x<y なら $ A_y円支給し、 $ x=y と更新する。
そうでないなら終了する。
青木さんの今月の給料は、この手順によって支給された金額の合計です。
青木さんの今月の給料の期待値を$ \mod 998244353 で求めてください。
/icons/hr.icon
制約
$ 1 \leq N \leq 3 \times 10^5
$ 0 \leq A_i < 998244353
/icons/hr.icon
出力
答えを出力せよ。
/icons/hr.icon
考察・解法
期待値 DP の基本的な問題だと思います。
期待値 DP は最近勉強していますが、個人的には最後から決めていき前に遡っていくイメージがあります。
今回は以下のような DP テーブルを考えます。
$ {\mathrm dp}[i] := 出目が i の時に貰える給料の期待値
このようにすると、求める答えは $ {\mathrm dp}[0] となります。
これを考えると $ x = Nの時は簡単です。
$ Nが出て$ A_N 円を獲得し、それ以降は何を出しても終わってしまうので、
$ {\mathrm dp}[N]=A_N です。
次に$ x = N - 1を考えます。この時、$ 1から$ N - 1が出てしまうと、その瞬間終わってしまうので給料は貰えませんが、$ x = Nが出ると追加で $ A_N円もらえます。
したがって、$ {\mathrm dp}[N - 1] = A_{N-1} + \dfrac{1}{N} \times A_N となります。
ここの $ \times A_Nというのが、$ x = Nの時に貰える期待値なので、
これは$ {\mathrm dp}[N] と考えることができます。
したがって、$ {\mathrm dp}[N - 1] = A_{N-1} + \dfrac{1}{N} \times {\mathrm dp}[N] となります。
次に$ x = N - 2について考えます。
$ 1から$ N - 2が出てしまうと、その瞬間終わってしまうので給料は貰えません。
それ以外で $ dp[N-1]+dp[N] 円が貰えます。これらの確率は何も$ x = N - 1と$ x = Nの場合なのでそれぞれ$ \dfrac{1}{N}です。
したがって、$ {\mathrm dp}[N - 2] = A_{N-2} + \dfrac{ {\mathrm dp}[N-1] + {\mathrm dp}[N] }{N} となります。
...
これを繰り返していくと、
$ x = kの時は
$ {\mathrm dp}[k] = A_k + \dfrac{1}{N} \sum_{i=0}^{k - 1} {\mathrm dp}[N-i]
これは後ろから求めていけばただの累積和になるので簡単に値を保持できます。このようにして後ろから辿っていけば$ {\mathrm dp}[0] を求めることが出来ました。
/icons/hr.icon
コード
code:python
n = int(input())
a = 0 + list(map(int, input().split())) mod = 998244353
p = pow(n, -1, mod)
num = 0
for i in reversed(range(n + 1)):
num %= mod
/icons/hr.icon
最後に
こちらの解説 を参考にさせていただきました。ありがとうございます。 期待値 dp やっぱり簡単なものでも難しいですね。まだまだ未熟です。